﻿/***********************/
/* (c) Alexandru Lungu */
/*  Challenge is Life! */
/* www.challenge-me.ws */
/***********************/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Coowork.WcfCommon;
using System.Xml;
using System.Security.Cryptography.Xml;
using System.Security.Cryptography;
using System.ServiceModel.Channels;
using System.IO;
using System.Security;
using ALungu.Security.Cryptography;
using System.Timers;
using System.Collections.Concurrent;


namespace Coowork.WcfServer
{
    //[9]
    public class ServerCryptographer : Cryptographer
    {
        //this should be biger than max of the credentials expires
        public const int AesBannedKeysExpiresTimeSpan = 10000;
        //for convenience the ban list will be clean at regulate intervals
        public const int AesBannedKeysCheckExpiresTimeSpan = 100000;

        static ConcurrentDictionary<string, DateTime> AesBannedKeys { get; set; }

        static Timer timer = new Timer();
        static ServerCryptographer()
        {
            AesBannedKeys = new ConcurrentDictionary<string, DateTime>();
            timer.Elapsed += new ElapsedEventHandler(timer_Elapsed);
        }

        static void timer_Elapsed(object sender, ElapsedEventArgs e)
        {
            string[] hashes = AesBannedKeys.Where(a => a.Value > DateTime.Now).Select(s => s.Key).ToArray();
            foreach (string hash in hashes)
            {
                DateTime banExpires;
                AesBannedKeys.TryRemove(hash, out banExpires);
            }
        }

        public static void Decrypt(XmlDocument xmlDoc)
        {
            if (xmlDoc == null)
                throw new ArgumentNullException("xmlDoc");

            XmlNodeList encryptedElements = xmlDoc.GetElementsByTagName("EncryptedData");
            if (encryptedElements.Count == 0)
                return;

            EncryptedXml encryptedXml = new EncryptedXml(xmlDoc);
            encryptedXml.AddKeyNameMapping(KeyElementName, RsaServiceProvider);

            EncryptedData eData = new EncryptedData();


            XmlElement encryptedElement = (XmlElement)encryptedElements[0];
            eData.LoadXml(encryptedElement);

            SymmetricAlgorithm a = encryptedXml.GetDecryptionKey(eData, eData.EncryptionMethod.KeyAlgorithm);



            if (AesKeys.ContainsKey(eData.Id))
            {
                throw new SecurityException("Id reuse inside the expiration timespan!");
            }
            else
            {
                //here the aes service provider gets the client key! [6]
                AesKeys.TryAdd(eData.Id, a.Key);
            }

            string keyHash = a.Key.ComputeHash();
            if (AesBannedKeys.ContainsKey(keyHash))
            {
                throw new SecurityException("Password reuse before ban expiration!");
            }
            else
            {
                AesBannedKeys.TryAdd(keyHash, DateTime.Now.AddMilliseconds(AesBannedKeysExpiresTimeSpan));
            }


            encryptedXml.DecryptDocument();
        }
        
        //[9]
        public static void Encrypt(XmlDocument xmlDoc, string elementToEncrypt)
        {
            if (xmlDoc == null)
                throw new ArgumentNullException("xmlDoc");

            XmlNodeList elementsToEncrypt = xmlDoc.GetElementsByTagName(elementToEncrypt);
            if (elementsToEncrypt.Count == 0)
                return;

            AesCryptoServiceProvider aesServiceProvider = new AesCryptoServiceProvider();

            XmlNode idNode = xmlDoc.GetElementsByTagName("a:RelatesTo")[0];
            string id = "";
            if (idNode != null)
            {
                id = idNode.InnerText;
            }
            
            byte[] key;
            AesKeys.TryRemove(id, out key);
            aesServiceProvider.Key = key;

            XmlElement xmlElementToEncrypt = (XmlElement)elementsToEncrypt[0];

            EncryptedXml encryptedXml = new EncryptedXml();
            byte[] encryptedElement = encryptedXml.EncryptData(xmlElementToEncrypt, aesServiceProvider, Content);

            EncryptedData encryptedData = new EncryptedData();
            encryptedData.Type = EncryptedXml.XmlEncElementUrl;
            encryptedData.EncryptionMethod = new EncryptionMethod(EncryptedXml.XmlEncAES256Url);
            encryptedData.CipherData.CipherValue = encryptedElement;
            encryptedData.Id = id;
            EncryptedXml.ReplaceElement(xmlElementToEncrypt, encryptedData, Content);
        }

        public static ArraySegment<byte> EncryptBuffer(ArraySegment<byte> buffer, BufferManager bufferManager, int messageOffset, string elementToEncrypt = "s:Envelope")
        {

            byte[] bufferedBytes;
            byte[] encryptedBytes;
            XmlDocument xmlDoc = new XmlDocument();

            using (MemoryStream memoryStream = new MemoryStream(buffer.Array, buffer.Offset, buffer.Count))
            {
                xmlDoc.Load(memoryStream);
            }
            ServerCryptographer.Encrypt(xmlDoc, elementToEncrypt);
            encryptedBytes = Encoding.UTF8.GetBytes(xmlDoc.OuterXml);
            bufferedBytes = bufferManager.TakeBuffer(encryptedBytes.Length);
            Array.Copy(encryptedBytes, 0, bufferedBytes, 0, encryptedBytes.Length);
            bufferManager.ReturnBuffer(buffer.Array);

            ArraySegment<byte> byteArray = new ArraySegment<byte>(bufferedBytes, messageOffset, encryptedBytes.Length);
            return byteArray;
        }

        public static ArraySegment<byte> DecryptBuffer(ArraySegment<byte> buffer, BufferManager bufferManager)
        {

            byte[] bufferedBytes;
            ArraySegment<byte> byteArray;
            XmlDocument xmlDoc = new XmlDocument();
            MemoryStream memoryStream = new MemoryStream(buffer.Array, buffer.Offset, buffer.Count);
            xmlDoc.Load(memoryStream);
            ServerCryptographer.Decrypt(xmlDoc);
            byte[] decryptedBytes = Encoding.UTF8.GetBytes(xmlDoc.OuterXml);

            bufferedBytes = bufferManager.TakeBuffer(decryptedBytes.Length);
            Array.Copy(decryptedBytes, 0, bufferedBytes, 0, decryptedBytes.Length);
            bufferManager.ReturnBuffer(buffer.Array);
            byteArray = new ArraySegment<byte>(bufferedBytes, 0, decryptedBytes.Length);
            return byteArray;
        }

    }

}

